package tools;

import java.util.*;

public class GrammarUtils {
    private final Grammar grammar;

    // FIRST集和FOLLOW集存储结构
    private final Map<String, Set<String>> firstSets = new HashMap<>();
    private final Map<String, Set<String>> followSets = new HashMap<>();

    // ε 用于表示空串

    //
    private static final String EPSILON = "ε";
    private static final String END_MARKER = "$";

    //构造函数
    public GrammarUtils(Grammar grammar) {
        this.grammar = grammar;
        initSets();
        computeFirstSets();
        computeFollowSets();
    }

    // 初始化 FIRST 和 FOLLOW 集合

    //
    private void initSets() {
        for (String nt : grammar.nonTerminals) {
            //对两个集合初始化
            firstSets.put(nt, new HashSet<>());
            followSets.put(nt, new HashSet<>());
        }
        for (String t : grammar.terminals) {
            //
            firstSets.put(t, new HashSet<>(Collections.singleton(t))); // 终结符的FIRST集是其自身
        }
        // 起始符FOLLOW加入$
        followSets.get(grammar.startSymbol).add(END_MARKER);
    }

    // 计算FIRST集
    private void computeFirstSets() {
        boolean changed;
        do {
            changed = false;
            for (Grammar.Production p : grammar.productions) {
                String lhs = p.lhs;
                List<String> rhs = p.rhs;
                Set<String> lhsFirst = firstSets.get(lhs);

                int beforeSize = lhsFirst.size();

                if (rhs.isEmpty()) {
                    // 产生空串
                    lhsFirst.add(EPSILON);
                } else {
                    boolean allNullable = true;
                    for (String symbol : rhs) {
                        Set<String> symbolFirst = firstSets.get(symbol);
                        if (symbolFirst == null) symbolFirst = new HashSet<>(); // 防止空指针
                        lhsFirst.addAll(withoutEpsilon(symbolFirst));
                        if (!symbolFirst.contains(EPSILON)) {
                            allNullable = false;
                            break;
                        }
                    }
                    if (allNullable) {
                        lhsFirst.add(EPSILON);
                    }
                }

                if (lhsFirst.size() > beforeSize) {
                    changed = true;
                }
            }
        } while (changed);
    }

    private Set<String> withoutEpsilon(Set<String> set) {
        Set<String> result = new HashSet<>(set);
        result.remove(EPSILON);
        return result;
    }

    // 计算FOLLOW集
    private void computeFollowSets() {
        boolean changed;
        do {
            changed = false;
            for (Grammar.Production p : grammar.productions) {
                String lhs = p.lhs;
                List<String> rhs = p.rhs;
                Set<String> lhsFollow = followSets.get(lhs);

                for (int i = 0; i < rhs.size(); i++) {
                    String B = rhs.get(i);
                    if (!grammar.nonTerminals.contains(B)) continue; // 只有非终结符才计算FOLLOW

                    Set<String> BFollow = followSets.get(B);
                    int beforeSize = BFollow.size();

                    // 计算β部分 FIRST集
                    Set<String> firstOfBeta = new HashSet<>();
                    boolean betaNullable = true;
                    for (int j = i + 1; j < rhs.size(); j++) {
                        String betaSymbol = rhs.get(j);
                        Set<String> betaFirst = firstSets.get(betaSymbol);
                        if (betaFirst == null) betaFirst = new HashSet<>();
                        firstOfBeta.addAll(withoutEpsilon(betaFirst));
                        if (!betaFirst.contains(EPSILON)) {
                            betaNullable = false;
                            break;
                        }
                    }

                    if (firstOfBeta.size() > 0) {
                        BFollow.addAll(firstOfBeta);
                    }

                    // 如果β可空，或者β为空，则把 FOLLOW(A) 加入 FOLLOW(B)
                    if (betaNullable) {
                        BFollow.addAll(lhsFollow);
                    }

                    if (BFollow.size() > beforeSize) {
                        changed = true;
                    }
                }
            }
        } while (changed);
    }

    public Map<String, Set<String>> getFirstSets() {
        return firstSets;
    }

    public Map<String, Set<String>> getFollowSets() {
        return followSets;
    }

    // 调试打印
    public void printSets() {
        System.out.println("=== FIRST Sets ===");
        for (String symbol : firstSets.keySet()) {
            System.out.println(symbol + ": " + firstSets.get(symbol));
        }
        System.out.println("=== FOLLOW Sets ===");
        for (String nt : grammar.nonTerminals) {
            System.out.println(nt + ": " + followSets.get(nt));
        }
    }
}
